{-# LANGUAGE RankNTypes,PatternGuards,NoMonomorphismRestriction,FlexibleInstances,UnboxedTuples #-}
{-# LANGUAGE ViewPatterns,DeriveGeneric,TemplateHaskell,ScopedTypeVariables,BangPatterns,FlexibleContexts #-}


module SliceSummary where

import SliceADT
import Utils
import ForwardIRSlice

import GHC.Generics ( Generic )
import Control.DeepSeq
import Control.DeepSeq.Generics ( genericRnf )
import Control.Lens ( Lens', makeLenses, (.~), (%~), (^.) )
import Control.Monad 
import Control.Monad.RWS.Strict

import Data.Map ( Map )
import qualified Data.Map as M
import Data.HashMap.Strict ( HashMap )
import qualified Data.HashMap.Strict as HM
import Data.IntMap ( IntMap )
import qualified Data.IntMap as IM
import Data.Set ( Set )
import qualified Data.Set as S
import Data.HashSet ( HashSet )
import qualified Data.HashSet as HS
import Data.IntSet ( IntSet )
import qualified Data.IntSet as IS
import Data.Maybe  
import Data.Monoid
import Data.List 
import Data.Char ( isLetter )
import qualified Data.Foldable as F
import qualified Data.Text as T (unpack)
import qualified Data.Text.Lazy as L ( unpack )
import Data.GraphViz hiding (Both)

import LLVM.Analysis
import LLVM.Analysis.CFG
import LLVM.Analysis.CallGraphSCCTraversal
import LLVM.Analysis.Dataflow
import LLVM.Analysis.PointsTo
import LLVM.Analysis.PointsTo.TrivialFunction
--import LLVM.Analysis.PointsTo.Andersen
import LLVM.Analysis.CallGraph

------
data SliceSummary =
  SliceSummary {  _bwdSliceTable :: Slices
                , _procTblSummary :: Map String Slices
                , _fwdSliceTable :: Slices
--                , _traceSize :: Int
                }
  deriving (Generic)

$(makeLenses ''SliceSummary)

instance Eq SliceSummary where
  (SliceSummary s1 fs1 d1) == (SliceSummary s2 fs2 d2) =
    s1 == s2 && fs1 == fs2 && d1 == d2

instance Monoid SliceSummary where
  mempty = SliceSummary mempty mempty mempty 
  mappend (SliceSummary bs1 ps1 fs1) (SliceSummary bs2 ps2 fs2) =
    SliceSummary (M.union bs2 bs1)(M.union ps1 ps2)(M.union fs2 fs1)

instance NFData SliceSummary where
  rnf = genericRnf

--
data SliceAnalysis = SliceAnalysis { _sliceSumm :: SliceSummary }
  deriving (Eq, Generic)

$(makeLenses ''SliceAnalysis)

instance NFData SliceAnalysis where
  rnf = genericRnf

instance Monoid SliceAnalysis where
  mempty = SliceAnalysis { _sliceSumm = mempty }
  mappend a1 a2 =
    SliceAnalysis { _sliceSumm = _sliceSumm a1 `mappend` _sliceSumm a2 }

---------------
----
data SliceEnv = SEnv { procSliTbl :: Map String Slices
                     , paraValMap :: IntMap (Value,Int)
                     }                      

data SliceState = SState { instDeps :: IntMap ValueIds
                         , traces :: !Int    -- [Instruction]
                         } deriving (Generic)


data SliceInfo = SInfo { _sliceTable :: !Slices
                       }
              deriving (Eq,Ord,Show,Generic)

$(makeLenses ''SliceInfo)

instance NFData SliceState where
  rnf = genericRnf

instance NFData SliceInfo where
  rnf = genericRnf

top :: SliceInfo
top = SInfo mempty 

type Analysis = RWS SliceEnv String SliceState 
type Transfer = SliceInfo -> Instruction -> Analysis SliceInfo

meetSliceInfo :: SliceInfo -> SliceInfo -> SliceInfo
meetSliceInfo (SInfo s1) (SInfo s2) = SInfo (mrgSli s1 s2)


updSliInfo :: Instruction -> (Value,Value) -> SliceInfo -> Analysis SliceInfo
updSliInfo i (ptr,v) si = do
  l' <- getUnionL' si i v
  let    
    baseName = toVarName ptr
    updOrXtdSli = if isAggregate ptr || isAggType ptr
                  then xtdSli else updSli
    si' = if isNothing baseName then si
           else (sliceTable %~ updOrXtdSli (fromJust baseName) l') si
  if null (refs ptr) 
   then addTrInstDep si i l'
   else case valueContent ptr of      
      InstructionC PhiNode {phiIncomingValues = (map fst -> ivs)} -> do
         let ptNames = mapMaybe toVarName ivs  
             si2 = if null ptNames then si 
                    else (sliceTable %~ xtdSli2 ptNames l') si    -- si'
         addTrInstDep si2 i l'
      InstructionC SelectInst {selectTrueValue = tv, selectFalseValue = fv} -> do
         let ptNames = mapMaybe toVarName [tv,fv]
             si2 = if null ptNames then si 
                    else (sliceTable %~ xtdSli2 ptNames l') si    -- si'
         addTrInstDep si2 i l'
      _ -> addTrInstDep si' i l'


identifySlice :: (FuncLike funcLike, HasFunction funcLike, HasCFG funcLike)
        => Module -> Lens' compositeSummary SliceSummary -> Transfer
        -> ComposableAnalysis compositeSummary funcLike
identifySlice m lns tr =
  composableAnalysisM runner (sliceAnalysis m tr) lns
  where
    runner a = fst (evalRWS a env st)
    env = SEnv mempty mempty   
    st = SState mempty 0 

sliceAnalysis :: (FuncLike funcLike, HasCFG funcLike,HasFunction funcLike)
               => Module -> Transfer -> funcLike -> SliceSummary -> Analysis SliceSummary
sliceAnalysis m tr funcLike s@(SliceSummary bs ptM fs)  = do  
  let envMod e = e { procSliTbl = ptM
                   , paraValMap = IM.fromList $ zip allParaIds allParaVals
                   }      
      frmlVals fn = zip (map toValue $ functionParameters fn) [0..]
      globalVals = zip (map toValue $ moduleGlobalVariables m) (repeat (-1))
--      allocVals fn = zip (map toValue $ funcAllocInsts fn) (repeat (-2))
      allParaVals = concatMap frmlVals (moduleDefinedFunctions m) ++ globalVals
--                    ++ concatMap allocVals (moduleDefinedFunctions m)  
      allParaIds = map (valueUniqueId . fst) allParaVals
      localParaVals = frmlVals f ++ globalVals  -- ++ allocVals f
      argStrs = map (toVarName' . fst) localParaVals
      argIds = map (IS.singleton . negate . valueUniqueId . fst) localParaVals
      fact0 = (sliceTable %~ updSlices argStrs argIds) top
--      fact0 = if findMain m == Just f then top else top'
      analysis = dataflowAnalysis top meetSliceInfo tr
  localInfo <- local envMod (forwardDataflow funcLike analysis fact0) 
--  trs <- getTrace 
  ideps <- getDeps        
  let fName = identifierAsString (functionName f)
      SInfo bs' = dataflowResult localInfo
      fs' = getFwdSlices (genValueMap m) ideps fs 
--  let traceStr = fName ++ fStrLn ++ "\'s #Insts_traced = " ++ show trs 
--      fStrLn = case getFunctionLine f of
--               Just (src, ln) -> "(Defined at " ++ src ++ ":" ++ show ln ++ ")"
--               Nothing        -> ""        
  return $ (procTblSummary %~ (M.insertWith' (\_ y->y) fName bs')) 
         $ (fwdSliceTable %~ M.unionWith IS.union fs') 
         $! (bwdSliceTable %~ M.union bs') s   -- `debug` traceStr
  where
    f = getFunction funcLike


modularSlice :: Module -> Transfer -> SliceSummary
modularSlice m tr = _sliceSumm res    -- `showGraph` (cg,mName)
  where
    cg = callGraph m pta []
    pta = runPointsToAnalysis m
--    mName = T.unpack $ moduleIdentifier m
    analyses :: [ComposableAnalysis SliceAnalysis Function]
    analyses = [ identifySlice m sliceSumm tr ]
    analysisFunc = callGraphComposeAnalysis analyses   
    res = callGraphSCCTraversal cg analysisFunc mempty    -- parallelCallGraphSCCTraversal

 
----------------------------------
------
getUnionL' :: IsValue a => SliceInfo -> Instruction -> a -> Analysis ValueIds
getUnionL' si i v = do
  l0 <- getInstDep i
  return $! unionL' l l0 (v,ss)
 where l = instructionUniqueId i
       ss = si ^. sliceTable
{-# INLINE getUnionL' #-}

addTrInstDep :: SliceInfo -> Instruction -> ValueIds -> Analysis SliceInfo
addTrInstDep si i l' = do 
  SState !deps trs <- get
  let deps' = IM.insertWith' IS.union (instructionUniqueId i) l' deps
      trs' = 1 + trs    -- i:trs
  put (SState deps' trs')
  return si 

setTrCtrDep :: SliceInfo -> Instruction -> ValueIds -> Analysis SliceInfo   
setTrCtrDep si i cis = do
  SState deps trs <- get
  l' <- getUnionL' si i i
  let deps' = IS.foldl' addF deps cis  
      addF m k = IM.insertWith' IS.union k l' m  
      trs' = 1 + trs    -- i:trs
--      cis' = IS.insert (instructionUniqueId i) cis
  put (SState deps' trs')
  return si

lkpInstDep ::  IntMap ValueIds -> Instruction -> ValueIds
lkpInstDep deps i = 
   IM.findWithDefault IS.empty (instructionUniqueId i) deps
   
getDeps :: Analysis (IntMap ValueIds)
getDeps = do {s <- get; return (instDeps s) }

getInstDep :: Instruction -> Analysis ValueIds
getInstDep i = do 
  deps <- getDeps  
  return $ lkpInstDep deps i 

addInstDep :: Instruction -> ValueIds -> Analysis ()
addInstDep i l' = do 
  SState !deps trs <- get  
  let deps' = IM.insertWith' IS.union (instructionUniqueId i) l' deps
  put (SState deps' trs)
  

getTrace :: Analysis Int    -- [Instruction]
getTrace = do {s <- get; return (traces s) }  

addTrace :: Instruction -> Analysis ()
addTrace i = do 
  s <- get  
  let tr' = 1 + traces s   -- -- i:(traces s)
  put s { traces = tr'}

----------------- 
--------
unionL' :: IsValue a => UniqueId -> ValueIds -> (a,Slices) -> ValueIds
unionL' l l0 (v,s) = IS.unions [IS.singleton l, l0, refIds, unionLkpSli refStrs s]
  where refValues = refVals v
        refIds = IS.fromList . HS.toList . HS.map valueUniqueId $ refValues 
        refStrs = mapMaybe toVarName . HS.toList $ refValues    -- refs
{-# INLINE unionL' #-}

   
